import secrets
import chipwhisperer as cw
import argparse, binascii, gc, gzip, matplotlib, matplotlib.pyplot as plt
import numpy, pickle, random, serial, sys, time
matplotlib.use( 'tkagg' )

# the tutorial of using CW305 with CW-Lite:
# https://github.com/newaetech/chipwhisperer-jupyter/blob/master/demos/PA_HW_CW305_1-Attacking_AES_on_an_FPGA.ipynb

# ===================================================================================================
# notes: the sample rate of CW-Lite is 105 MS/s
# given the core frequency is 8 MHz, this means 13.125 samples per cycle

# connect and config the scope (i.e., CW-Lite)
scope = cw.scope()                            # connect to the CW-Lite
scope.adc.samples = 300                       # #samples after  triggering (~60 cc)
scope.adc.presamples = 0	                  # #samples before triggering
scope.adc.offset = 0                          # triggering offset
scope.adc.basic_mode = "rising_edge"          # rising edge to trigger
scope.gain.db = 20                            # low-noise amplifier (0 - 55 db)
# target board info
TARGET_PLATFORM = 'CW305_100t'                # the target is CW305-100t
fpga_id = '100t'                              # the FPGA type is 100t
# connect the target board with scope
bitstream = r"Path_to_bitstream"
target = cw.target(
    scope,                                    # scope (i.e., CW-Lite)
    cw.targets.CW305,                         # the target board
    fpga_id = fpga_id,                        # FPGA type
    bsfile = bitstream,                       # the bitstream file
    force = True)                             # needs to erase the board manually
# config the target board
target.vccint_set(1.0)                        # configure VCC
target.pll.pll_enable_set(True)               # enable  PLL
target.pll.pll_outenable_set(False, 0)        # disable PLL0
target.pll.pll_outenable_set(True, 1)         # enable  PLL1
target.pll.pll_outenable_set(False, 2)        # disable PLL2
target.pll.pll_outfreq_set(16E6, 1)           # PLL frequency is 16MHz
target.clkusbautooff = True                   # configure USB clock
# ensure the ADC clock is locked
for i in range(5):
    scope.clock.reset_adc()
    time.sleep(1)
    if scope.clock.adc_locked:
        break
assert (scope.clock.adc_locked), "ADC failed to lock"
# ===================================================================================================

# format conversions
def bytes2seq( x ) :
  return            [ int( t ) for t in x ]

def seq2bytes( x ) :
  return bytearray( [ int( t ) for t in x ] )

def octetstr2bytes( x ) :
  t = x.split( ':' ) ; n = int( t[ 0 ], 16 ) ; x = binascii.a2b_hex( t[ 1 ] )
  if ( n != len( x ) ) :
    raise ValueError
  else :
    return x

def bytes2octetstr( x ) :
  n = '{0:02X}'.format( len( x ) )
  x = binascii.b2a_hex( x ).decode( 'ascii' ).upper()
  return n + ':' + x

def randbytes( n, seed = None ) :
  if ( seed != None ) :
    t = random.Random( seed )
  else :
    t = random
  return bytes( [ t.getrandbits( 8 ) for i in range( n ) ] )

def resize( xs, n, dtype = numpy.float32 ) :
  if   ( len( xs ) <  n ) :
    return numpy.concatenate( ( xs[ 0 :   ], numpy.array( [ 0 ] * ( n - len( xs ) ), dtype = dtype ) ) )
  elif ( len( xs ) >  n ) :
    return                      xs[ 0 : n ]
  elif ( len( xs ) == n ) :
    return                      xs
# =============================================================================

class Board( object ) :
  def __init__( self ) :
    pass

  def  open( self ) :
    self.board = serial.Serial( 
      port = argv.uart_port, 
      baudrate = argv.uart_baud,
      bytesize = serial.EIGHTBITS, 
      parity = serial.PARITY_NONE, 
      stopbits = serial.STOPBITS_ONE, 
      timeout = None )
  
  def close( self ) :
    self.board.close()

  def reset( self ) :
    self.board.reset_input_buffer()
    self.board.reset_output_buffer()

  def ln_wr( self, x ) :
    self.board.write((x + '\x0D').encode('ascii'))
    self.board.flush()
    # time.sleep( 0.1 )

  def ln_rd( self    ) :
    r = ''
    while( True ):
      t = self.board.read( 1 ).decode( 'ascii' )
      if ( t == '\x0D' ) :
        break
      else :
        r += t
    # time.sleep( 0.1 )
    return r

# -----------------------------------------------------------------------------

def acquire() :
    board = Board()
    board.open()

    # for batch acquisition
    batch_ctr = 0
    batch_size = 5000

    A = numpy.zeros(batch_size, dtype=numpy.uint32)

    for i in range( argv.dataset_size ) :
        if ((i % batch_size) == 0): print( 'acquire> iteration {0:d}/{1:d}'.format( i, argv.dataset_size) )
        while (True):
            # ISW
            bytearr = secrets.token_bytes(12)
            x0_r = bytearr[0:2] + b'\x00' + bytearr[2:3]
            x1_r = bytearr[3:5] + b'\x00' + bytearr[5:6]
            y0_r = bytearr[6:8] + b'\x00' + bytearr[8:9]
            y1_r = bytearr[9:11] + b'\x00' + bytearr[11:12]

            # AES S-box
            # bytearr = secrets.token_bytes(24)
            # x0_r = bytearr[0:2] + b'\x00' + bytearr[2:3]
            # x1_r = bytearr[3:5] + b'\x00' + bytearr[5:6]
            # x2_r = bytearr[6:8] + b'\x00' + bytearr[8:9]
            # x3_r = bytearr[9:11] + b'\x00' + bytearr[11:12]
            # x4_r = bytearr[12:14] + b'\x00' + bytearr[14:15]
            # x5_r = bytearr[15:17] + b'\x00' + bytearr[17:18]
            # x6_r = bytearr[18:20] + b'\x00' + bytearr[20:21]
            # x7_r = bytearr[21:23] + b'\x00' + bytearr[23:24]
            # bytearr2 = secrets.token_bytes(24)
            # y0_r = bytearr2[0:2] + b'\x00' + bytearr2[2:3]
            # y1_r = bytearr2[3:5] + b'\x00' + bytearr2[5:6]
            # y2_r = bytearr2[6:8] + b'\x00' + bytearr2[8:9]
            # y3_r = bytearr2[9:11] + b'\x00' + bytearr2[11:12]
            # y4_r = bytearr2[12:14] + b'\x00' + bytearr2[14:15]
            # y5_r = bytearr2[15:17] + b'\x00' + bytearr2[17:18]
            # y6_r = bytearr2[18:20] + b'\x00' + bytearr2[20:21]
            # y7_r = bytearr2[21:23] + b'\x00' + bytearr2[23:24]

            # arm the scope
            scope.arm()

            # send data
            # ISW
            board.ln_wr( bytes2octetstr( x0_r ) )
            board.ln_wr( bytes2octetstr( x1_r ) )
            board.ln_wr( bytes2octetstr( y0_r ) )
            board.ln_wr( bytes2octetstr( y1_r ) )

            # AES S-box
            # board.ln_wr(bytes2octetstr(x0_r))
            # board.ln_wr(bytes2octetstr(x1_r))
            # board.ln_wr(bytes2octetstr(x2_r))
            # board.ln_wr(bytes2octetstr(x3_r))
            # board.ln_wr(bytes2octetstr(x4_r))
            # board.ln_wr(bytes2octetstr(x5_r))
            # board.ln_wr(bytes2octetstr(x6_r))
            # board.ln_wr(bytes2octetstr(x7_r))
            # board.ln_wr(bytes2octetstr(y0_r))
            # board.ln_wr(bytes2octetstr(y1_r))
            # board.ln_wr(bytes2octetstr(y2_r))
            # board.ln_wr(bytes2octetstr(y3_r))
            # board.ln_wr(bytes2octetstr(y4_r))
            # board.ln_wr(bytes2octetstr(y5_r))
            # board.ln_wr(bytes2octetstr(y6_r))
            # board.ln_wr(bytes2octetstr(y7_r))

            flag = bytes2octetstr( x0_r )

            # r_r  = octetstr2bytes( board.ln_rd() )
            # print( 'acquire> r_r = {0:s}'.format( bytes2octetstr( r_r ) ) )

            # capture the trace
            ret = scope.capture()

            break

        # return the captured trace
        r = scope.get_last_trace()
        
        if (i == 0): 
            R = numpy.ndarray( shape = ( batch_size, len( r ) ), dtype = numpy.float32 )
            len_r0 = len(r)

        resize( r, R.shape[ 1 ], dtype = numpy.float32 )

        for ( j, x ) in enumerate( r[ 0 : R.shape[ 1 ] ] ) :
            R[ i-batch_ctr*batch_size, j ] = R.dtype.type( x )

        flag = '0x' + flag[3:len(flag)]
        flag = int(flag,16)
        A[i-batch_ctr*batch_size] = flag

        del r; gc.collect()

        if (((i+1) % batch_size) == 0):
            fd = open( "{}-{}".format(argv.dataset_wr,batch_ctr), 'wb' )

            pickle.dump( R, fd )
            pickle.dump( A, fd )

            fd.close()

            numpy.save("E:/CBM/traces/snr-base-core-cbm-isw/trace-batch/" + "tracebatch-{}.npy".format(batch_ctr), R)
            numpy.save("E:/CBM/traces/snr-base-core-cbm-isw/x0-batch/" + "x0batch-{}.npy".format(batch_ctr), A)

            del R,A,flag

            batch_ctr += 1

            A = numpy.zeros(batch_size, dtype=numpy.uint32)
            R = numpy.ndarray( shape = ( batch_size, len_r0 ), dtype = numpy.float32 )

    # disconnect the scope and board
    board.close()
    scope.dis()
    target.dis()

    return R

# =============================================================================

if ( __name__ == '__main__' ) :
  parser = argparse.ArgumentParser()
  parser.add_argument( '--uart-port',    dest = 'uart_port',    type = str, action = 'store', default = 'COM4' )
  parser.add_argument( '--uart-baud',    dest = 'uart_baud',    type = int, action = 'store', default = 115200  )
  parser.add_argument( '--dataset-seed', dest = 'dataset_seed', type = int, action = 'store', default = 0  )
  parser.add_argument( '--dataset-size', dest = 'dataset_size', type = int, action = 'store', default = 300000  )
  parser.add_argument( '--dataset-wr',   dest = 'dataset_wr',   type = str, action = 'store', default = 'E:/CBM/traces/snr-base-core-cbm-isw/batch/tr-sw-base-batch' )
  argv = parser.parse_args()
  R = acquire()
  print( 'acquired |R| = {0:d} x {1:d}'.format( R.shape[ 0 ], R.shape[ 1 ]   ) )

